{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Jxv6goXm7oGF"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "llMNufAK7nfK"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8Byow2J6LaPl"
      },
      "source": [
        "# Better performance with tf.function and AutoGraph"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kGXS3UWBBNoc"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/guide/function\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/function.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/function.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/function.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "CydFK2CL7ZHA"
      },
      "source": [
        "TF 2.0 brings together the ease of eager execution and the power of TF 1.0. At the center of this merger is `tf.function`, which allows you to transform a subset of Python syntax into portable, high-performance TensorFlow graphs.\n",
        "\n",
        "A cool new feature of `tf.function` is AutoGraph, which lets you write graph code using natural Python syntax. For a list of the Python features that you can use with AutoGraph, see [AutoGraph Capabilities and Limitations](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/limitations.md). For more details about `tf.function`, see the RFC [TF 2.0: Functions, not Sessions](https://github.com/tensorflow/community/blob/master/rfcs/20180918-functions-not-sessions-20.md). For more details about AutoGraph, see `tf.autograph`.\n",
        "\n",
        "This tutorial will walk you through the basic features of `tf.function` and AutoGraph."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "n4EKOpw9mObL"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Import TensorFlow 2.0:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "V9oECvVSI1Kj"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import, division, print_function, unicode_literals\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mT7meGqrZTz9"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  # %tensorflow_version only exists in Colab.\n",
        "  %tensorflow_version 2.x\n",
        "except Exception:\n",
        "  pass\n",
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "77AsVr1GGtBP"
      },
      "source": [
        "## The `tf.function` decorator\n",
        "\n",
        "When you annotate a function with `tf.function`, you can still call it like any other function. But it will be compiled into a graph, which means you get the benefits of faster execution, running on GPU or TPU, or exporting to SavedModel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "FhIg7-z6HNWj"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def simple_nn_layer(x, y):\n",
        "  return tf.nn.relu(tf.matmul(x, y))\n",
        "\n",
        "\n",
        "x = tf.random.uniform((3, 3))\n",
        "y = tf.random.uniform((3, 3))\n",
        "\n",
        "simple_nn_layer(x, y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "U-LAE4pMNR9g"
      },
      "source": [
        "If we examine the result of the annotation, we can see that it's a special callable that handles all interactions with the TensorFlow runtime."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "q4t2iuS7Nqc0"
      },
      "outputs": [],
      "source": [
        "simple_nn_layer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DqeefLGNXjZQ"
      },
      "source": [
        "If your code uses multiple functions, you don't need to annotate them all - any functions called from an annotated function will also run in graph mode."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3VGF7tlVXiZY"
      },
      "outputs": [],
      "source": [
        "def linear_layer(x):\n",
        "  return 2 * x + 1\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def deep_net(x):\n",
        "  return tf.nn.relu(linear_layer(x))\n",
        "\n",
        "\n",
        "deep_net(tf.constant((1, 2, 3)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yQvg6ZSKWyqE"
      },
      "source": [
        "Functions can be faster than eager code, for graphs with many small ops. But for graphs with a few expensive ops (like convolutions), you may not see much speedup.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "0EL6lVwEWuFo"
      },
      "outputs": [],
      "source": [
        "import timeit\n",
        "conv_layer = tf.keras.layers.Conv2D(100, 3)\n",
        "\n",
        "@tf.function\n",
        "def conv_fn(image):\n",
        "  return conv_layer(image)\n",
        "\n",
        "image = tf.zeros([1, 200, 200, 100])\n",
        "# warm up\n",
        "conv_layer(image); conv_fn(image)\n",
        "print(\"Eager conv:\", timeit.timeit(lambda: conv_layer(image), number=10))\n",
        "print(\"Function conv:\", timeit.timeit(lambda: conv_fn(image), number=10))\n",
        "print(\"Note how there's not much difference in performance for convolutions\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "L4zj-jpH0jKH"
      },
      "outputs": [],
      "source": [
        "lstm_cell = tf.keras.layers.LSTMCell(10)\n",
        "\n",
        "@tf.function\n",
        "def lstm_fn(input, state):\n",
        "  return lstm_cell(input, state)\n",
        "\n",
        "input = tf.zeros([10, 10])\n",
        "state = [tf.zeros([10, 10])] * 2\n",
        "# warm up\n",
        "lstm_cell(input, state); lstm_fn(input, state)\n",
        "print(\"eager lstm:\", timeit.timeit(lambda: lstm_cell(input, state), number=10))\n",
        "print(\"function lstm:\", timeit.timeit(lambda: lstm_fn(input, state), number=10))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ohbSnA79mcJV"
      },
      "source": [
        "## Use Python control flow\n",
        "\n",
        "When using data-dependent control flow inside `tf.function`, you can use Python control flow statements and AutoGraph will convert them into appropriate TensorFlow ops. For example, `if` statements will be converted into `tf.cond()` if they depend on a `Tensor`.\n",
        "\n",
        "In the example below, `x` is a `Tensor` but the `if` statement works as expected:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aA3gOodCBkOw"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def square_if_positive(x):\n",
        "  if x \u003e 0:\n",
        "    x = x * x\n",
        "  else:\n",
        "    x = 0\n",
        "  return x\n",
        "\n",
        "\n",
        "print('square_if_positive(2) = {}'.format(square_if_positive(tf.constant(2))))\n",
        "print('square_if_positive(-2) = {}'.format(square_if_positive(tf.constant(-2))))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "GMiCUkdyoq98"
      },
      "source": [
        "Note: The previous example uses simple conditionals with scalar values. \u003ca href=\"#batching\"\u003eBatching\u003c/a\u003e is typically used in real-world code."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "m-jWmsCmByyw"
      },
      "source": [
        "AutoGraph supports common Python statements like `while`, `for`, `if`, `break`, `continue` and `return`, with support for nesting. That means you can use `Tensor` expressions in the condition of `while` and `if` statements, or iterate over a `Tensor` in a `for` loop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "toxKBOXbB1ro"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def sum_even(items):\n",
        "  s = 0\n",
        "  for c in items:\n",
        "    if c % 2 \u003e 0:\n",
        "      continue\n",
        "    s += c\n",
        "  return s\n",
        "\n",
        "\n",
        "sum_even(tf.constant([10, 12, 15, 20]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "AtDaLrbySw4j"
      },
      "source": [
        "AutoGraph also provides a low-level API for advanced users. For example we can use it to have a look at the generated code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aRsde3x_SjTQ"
      },
      "outputs": [],
      "source": [
        "print(tf.autograph.to_code(sum_even.python_function))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "rvJXCfk8VkLf"
      },
      "source": [
        "Here's an example of more complicated control flow:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "h-Z87IJqVlKl"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def fizzbuzz(n):\n",
        "  for i in tf.range(n):\n",
        "    if i % 3 == 0:\n",
        "      tf.print('Fizz')\n",
        "    elif i % 5 == 0:\n",
        "      tf.print('Buzz')\n",
        "    else:\n",
        "      tf.print(i)\n",
        "\n",
        "fizzbuzz(tf.constant(15))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "h_Y4uC1R1B55"
      },
      "source": [
        "## Keras and AutoGraph\n",
        "\n",
        "AutoGraph is available by default in non-dynamic Keras models. For more information, see `tf.keras`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cR6mpLKP1HLe"
      },
      "outputs": [],
      "source": [
        "class CustomModel(tf.keras.models.Model):\n",
        "\n",
        "  @tf.function\n",
        "  def call(self, input_data):\n",
        "    if tf.reduce_mean(input_data) \u003e 0:\n",
        "      return input_data\n",
        "    else:\n",
        "      return input_data // 2\n",
        "\n",
        "\n",
        "model = CustomModel()\n",
        "\n",
        "model(tf.constant([-2, -4]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NTEvpBK9f8kj"
      },
      "source": [
        "## Side effects\n",
        "\n",
        "Just like in eager mode, you can use operations with side effects, like `tf.assign` or `tf.print` normally inside `tf.function`, and it will insert the necessary control dependencies to ensure they execute in order."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-Wd6i8S9gcuC"
      },
      "outputs": [],
      "source": [
        "v = tf.Variable(5)\n",
        "\n",
        "@tf.function\n",
        "def find_next_odd():\n",
        "  v.assign(v + 1)\n",
        "  if v % 2 == 0:\n",
        "    v.assign(v + 1)\n",
        "\n",
        "\n",
        "find_next_odd()\n",
        "v"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4LfnJjm0Bm0B"
      },
      "source": [
        "\u003ca id=\"debugging\"\u003e\u003c/a\u003e\n",
        "\n",
        "## Debugging\n",
        "\n",
        "`tf.function` and AutoGraph work by generating code and tracing it into TensorFlow graphs. This mechanism does not yet support step-by-step debuggers like `pdb`. However, you can call `tf.config.experimental_run_functions_eagerly(True)` to temporarily enable eager execution inside the `tf.function' and use your favorite debugger:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Yci8ve6hmgpF"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def f(x):\n",
        "  if x \u003e 0:\n",
        "    # Try setting a breakpoint here!\n",
        "    # Example:\n",
        "    #   import pdb\n",
        "    #   pdb.set_trace()\n",
        "    x = x + 1\n",
        "  return x\n",
        "\n",
        "tf.config.experimental_run_functions_eagerly(True)\n",
        "\n",
        "# You can now set breakpoints and run the code in a debugger.\n",
        "f(tf.constant(1))\n",
        "\n",
        "tf.config.experimental_run_functions_eagerly(False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NfpIQUv28Ht4"
      },
      "source": [
        "## Advanced example: An in-graph training loop\n",
        "\n",
        "The previous section showed that AutoGraph can be used inside Keras layers and models. Keras models can also be used in AutoGraph code.\n",
        "\n",
        "This example shows how to train a simple Keras model on MNIST with the entire training process—loading batches, calculating gradients, updating parameters, calculating validation accuracy, and repeating until convergence—is performed in-graph."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Em5dzSUOtLRP"
      },
      "source": [
        "### Download data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "xqoxumv0ssQW"
      },
      "outputs": [],
      "source": [
        "def prepare_mnist_features_and_labels(x, y):\n",
        "  x = tf.cast(x, tf.float32) / 255.0\n",
        "  y = tf.cast(y, tf.int64)\n",
        "  return x, y\n",
        "\n",
        "def mnist_dataset():\n",
        "  (x, y), _ = tf.keras.datasets.mnist.load_data()\n",
        "  ds = tf.data.Dataset.from_tensor_slices((x, y))\n",
        "  ds = ds.map(prepare_mnist_features_and_labels)\n",
        "  ds = ds.take(20000).shuffle(20000).batch(100)\n",
        "  return ds\n",
        "\n",
        "train_dataset = mnist_dataset()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "znmy4l8ntMvW"
      },
      "source": [
        "### Define the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ltxyJVWTqNAO"
      },
      "outputs": [],
      "source": [
        "model = tf.keras.Sequential((\n",
        "    tf.keras.layers.Reshape(target_shape=(28 * 28,), input_shape=(28, 28)),\n",
        "    tf.keras.layers.Dense(100, activation='relu'),\n",
        "    tf.keras.layers.Dense(100, activation='relu'),\n",
        "    tf.keras.layers.Dense(10)))\n",
        "model.build()\n",
        "optimizer = tf.keras.optimizers.Adam()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "oeYV6mKnJGMr"
      },
      "source": [
        "### Define the training loop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3xtg_MMhJETd"
      },
      "outputs": [],
      "source": [
        "compute_loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
        "\n",
        "compute_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()\n",
        "\n",
        "\n",
        "def train_one_step(model, optimizer, x, y):\n",
        "  with tf.GradientTape() as tape:\n",
        "    logits = model(x)\n",
        "    loss = compute_loss(y, logits)\n",
        "\n",
        "  grads = tape.gradient(loss, model.trainable_variables)\n",
        "  optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
        "\n",
        "  compute_accuracy(y, logits)\n",
        "  return loss\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def train(model, optimizer):\n",
        "  train_ds = mnist_dataset()\n",
        "  step = 0\n",
        "  loss = 0.0\n",
        "  accuracy = 0.0\n",
        "  for x, y in train_ds:\n",
        "    step += 1\n",
        "    loss = train_one_step(model, optimizer, x, y)\n",
        "    if step % 10 == 0:\n",
        "      tf.print('Step', step, ': loss', loss, '; accuracy', compute_accuracy.result())\n",
        "  return step, loss, accuracy\n",
        "\n",
        "step, loss, accuracy = train(model, optimizer)\n",
        "print('Final step', step, ': loss', loss, '; accuracy', compute_accuracy.result())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SnsumiP6eRYL"
      },
      "source": [
        "## Batching\n",
        "\n",
        "In real applications batching is essential for performance. The best code to convert to AutoGraph is code where the control flow is decided at the _batch_ level. If making decisions at the individual _example_ level, try to use batch APIs to maintain performance.\n",
        "\n",
        "For example, if you have the following code in Python:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "t31QoERiNccJ"
      },
      "outputs": [],
      "source": [
        "def square_if_positive(x):\n",
        "  return [i ** 2 if i \u003e 0 else i for i in x]\n",
        "\n",
        "\n",
        "square_if_positive(range(-5, 5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kSeEJ76uNgwD"
      },
      "source": [
        "You may be tempted to write it in TensorFlow as such (and this would work!):\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RqR8WzSzNf87"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def square_if_positive_naive(x):\n",
        "  result = tf.TensorArray(tf.int32, size=x.shape[0])\n",
        "  for i in tf.range(x.shape[0]):\n",
        "    if x[i] \u003e 0:\n",
        "      result = result.write(i, x[i] ** 2)\n",
        "    else:\n",
        "      result = result.write(i, x[i])\n",
        "  return result.stack()\n",
        "\n",
        "\n",
        "square_if_positive_naive(tf.range(-5, 5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gTcyWXVGN3gS"
      },
      "source": [
        "But in this case, it turns out you can write the following:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VO2f6x-lNfVj"
      },
      "outputs": [],
      "source": [
        "def square_if_positive_vectorized(x):\n",
        "  return tf.where(x \u003e 0, x ** 2, x)\n",
        "\n",
        "\n",
        "square_if_positive_vectorized(tf.range(-5, 5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "EXFVx-cJ57_F"
      },
      "source": [
        "## Re-tracing\n",
        "\n",
        "Key points:\n",
        "\n",
        "* Exercise caution when calling functions with non-tensor arguments, or with arguments that change shapes.\n",
        "* Decorate module-level functions, and methods of module-level classes, and avoid decorating local functions or methods.\n",
        "\n",
        "`tf.function` can give you significant speedup over eager execution, at the cost of a slower first-time execution. This is because when executed for the first time, the function is also *traced* into a TensorFlow graph. Constructing and optimizing a graph is usually much slower compared to actually executing it:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "iahT-4wT6vlA"
      },
      "outputs": [],
      "source": [
        "import timeit\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def f(x, y):\n",
        "  return tf.matmul(x, y)\n",
        "\n",
        "print(\n",
        "    \"First invocation:\",\n",
        "    timeit.timeit(lambda: f(tf.ones((10, 10)), tf.ones((10, 10))), number=1))\n",
        "\n",
        "print(\n",
        "    \"Second invocation:\",\n",
        "    timeit.timeit(lambda: f(tf.ones((10, 10)), tf.ones((10, 10))), number=1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "q0Wojo2Z7hKg"
      },
      "source": [
        "You can easily tell when a function is traced by adding a `print` statement to the top of the function. Because any Python code is only executed at trace time, you will only see the output of `print` when the function is traced:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "2IHE7-jT7gZs"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def f():\n",
        "  print('Tracing!')\n",
        "  tf.print('Executing')\n",
        "\n",
        "print('First invocation:')\n",
        "f()\n",
        "\n",
        "print('Second invocation:')\n",
        "f()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "J2VBoQC58PdU"
      },
      "source": [
        "`tf.function` may also *re-trace* when called with different non-tensor arguments:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-c6VUwrz808l"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def f(n):\n",
        "  print(n, 'Tracing!')\n",
        "  tf.print(n, 'Executing')\n",
        "\n",
        "f(1)\n",
        "f(1)\n",
        "\n",
        "f(2)\n",
        "f(2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "aKOrjBLCE8cy"
      },
      "source": [
        "A *re-trace* can also happen when tensor arguments change shape, unless you specified an `input_signature`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "txhtkn0rE8dH"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def f(x):\n",
        "  print(x.shape, 'Tracing!')\n",
        "  tf.print(x, 'Executing')\n",
        "\n",
        "f(tf.constant([1]))\n",
        "f(tf.constant([2]))\n",
        "\n",
        "f(tf.constant([1, 2]))\n",
        "f(tf.constant([3, 4]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "sdN40ZqT9XaG"
      },
      "source": [
        "In addition, tf.function always creates a new graph function with its own set of traces whenever it is called:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "GT1iBa5i9enE"
      },
      "outputs": [],
      "source": [
        "def f():\n",
        "  print('Tracing!')\n",
        "  tf.print('Executing')\n",
        "\n",
        "tf.function(f)()\n",
        "tf.function(f)()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "n7_JDzFK9nnC"
      },
      "source": [
        "This can lead to surprising behavior when using the `@tf.function` decorator in a nested function:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "P3pBG7Uf9u4g"
      },
      "outputs": [],
      "source": [
        "def outer():\n",
        "  @tf.function\n",
        "  def f():\n",
        "    print('Tracing!')\n",
        "    tf.print('Executing')\n",
        "  f()\n",
        "\n",
        "outer()\n",
        "outer()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "Jxv6goXm7oGF"
      ],
      "name": "function.ipynb",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1Hl4PR32Y_nhq4AZI-ku9FlpZIgXjEU7d",
          "timestamp": 1544141357531
        },
        {
          "file_id": "https://github.com/tensorflow/docs/blob/master/site/en/guide/function.ipynb",
          "timestamp": 1542126584672
        }
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
